import os
from fastapi import FastAPI, UploadFile, File, HTTPException
from fastapi.responses import FileResponse
from fastapi.staticfiles import StaticFiles
from fastapi.responses import HTMLResponse
from typing import List
import shutil
import uuid
from datetime import datetime
import pytesseract
from PIL import Image
import numpy as np
import sys

# 设置Tesseract数据路径
if sys.platform == 'darwin':  # macOS
    tessdata_path = '/opt/homebrew/share/tessdata'
    if os.path.exists(tessdata_path):
        os.environ['TESSDATA_PREFIX'] = tessdata_path
    else:
        # 尝试其他可能的路径
        alt_paths = [
            '/usr/local/share/tessdata',
            '/usr/share/tessdata'
        ]
        for path in alt_paths:
            if os.path.exists(path):
                os.environ['TESSDATA_PREFIX'] = path
                break

# 尝试导入PaddleOCR
try:
    from paddleocr import PaddleOCR
    PADDLE_AVAILABLE = True
except ImportError:
    PADDLE_AVAILABLE = False
    print("警告: PaddleOCR未安装，将使用Tesseract进行OCR", file=sys.stderr)

app = FastAPI(
    title="文档处理API",
    description="提供PDF和图片处理功能的RESTful API",
    version="1.0.0"
)

# 创建临时文件目录
TEMP_DIR = "temp"
os.makedirs(TEMP_DIR, exist_ok=True)

# 挂载静态文件目录
app.mount("/static", StaticFiles(directory="static"), name="static")

@app.get("/", response_class=HTMLResponse)
async def get_index():
    """返回主页"""
    with open("static/index.html", "r", encoding="utf-8") as f:
        return f.read()

def get_temp_path(filename: str) -> str:
    """生成临时文件路径"""
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    unique_id = str(uuid.uuid4())[:8]
    return os.path.join(TEMP_DIR, f"{timestamp}_{unique_id}_{filename}")

def preprocess_image(image: Image.Image) -> Image.Image:
    """预处理图片以提高OCR准确率"""
    # 转换为灰度图
    image = image.convert('L')
    
    # 增加对比度
    from PIL import ImageEnhance
    enhancer = ImageEnhance.Contrast(image)
    image = enhancer.enhance(2.0)
    
    return image

def check_tesseract_installation() -> str:
    """检查Tesseract是否正确安装和配置"""
    try:
        # 检查tesseract命令是否可用
        pytesseract.get_tesseract_version()
        
        # 检查语言包
        lang = 'chi_sim'  # 简体中文
        if not os.path.exists(os.path.join(os.environ.get('TESSDATA_PREFIX', ''), f'{lang}.traineddata')):
            print(f"警告: 未找到中文语言包 ({lang}.traineddata)，将使用英文", file=sys.stderr)
            return 'eng'
        return lang
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Tesseract未正确安装或配置: {str(e)}")

@app.post("/api/images/to-pdf")
async def images_to_pdf(files: List[UploadFile] = File(...), dpi: int = 200):
    """
    将多个图片合并成一个PDF文件
    
    Args:
        files: 上传的图片文件列表
        dpi: 图片DPI（默认：200）
    
    Returns:
        生成的PDF文件
    """
    try:
        # 保存上传的图片
        image_paths = []
        for file in files:
            temp_path = get_temp_path(file.filename)
            with open(temp_path, "wb") as buffer:
                shutil.copyfileobj(file.file, buffer)
            image_paths.append(temp_path)
        
        # 创建输出PDF路径
        output_pdf = get_temp_path("output.pdf")
        
        # 打开第一张图片
        first_image = Image.open(image_paths[0]).convert('RGB')
        images = []
        
        # 转换其他图片
        for img_path in image_paths[1:]:
            img = Image.open(img_path).convert('RGB')
            images.append(img)
        
        # 保存为PDF
        first_image.save(output_pdf, save_all=True, append_images=images, dpi=(dpi, dpi))
        
        # 清理临时图片文件
        for path in image_paths:
            os.remove(path)
        
        return FileResponse(
            output_pdf,
            media_type="application/pdf",
            filename="output.pdf",
            background=None
        )
        
    except Exception as e:
        # 清理临时文件
        for path in image_paths:
            if os.path.exists(path):
                os.remove(path)
        if os.path.exists(output_pdf):
            os.remove(output_pdf)
        raise HTTPException(status_code=500, detail=str(e))

@app.post("/api/images/to-text")
async def image_to_text(
    file: UploadFile = File(...),
    lang: str = "ch",
    use_paddle: bool = True
):
    """
    将图片转换为文字
    
    Args:
        file: 上传的图片文件
        lang: OCR语言（默认：ch）
        use_paddle: 是否使用PaddleOCR（默认：True）
    
    Returns:
        识别的文本内容
    """
    try:
        # 保存上传的图片
        temp_path = get_temp_path(file.filename)
        with open(temp_path, "wb") as buffer:
            shutil.copyfileobj(file.file, buffer)
        
        # 打开图片
        image = Image.open(temp_path)
        
        # 如果PaddleOCR不可用，强制使用Tesseract
        if use_paddle and not PADDLE_AVAILABLE:
            use_paddle = False
        
        if use_paddle:
            # 使用PaddleOCR
            ocr = PaddleOCR(use_angle_cls=True, lang=lang, show_log=False)
            result = ocr.ocr(np.array(image), cls=True)
            
            # 提取文本
            text = ""
            if result:
                for line in result:
                    for item in line:
                        text += item[1][0] + "\n"
        else:
            # 使用Tesseract
            # 检查Tesseract安装
            tesseract_lang = check_tesseract_installation()
            
            # 预处理图片
            processed_image = preprocess_image(image)
            
            # 使用pytesseract进行OCR
            text = pytesseract.image_to_string(processed_image, lang=tesseract_lang)
            
            # 如果文本为空，尝试使用原始图片
            if not text.strip():
                text = pytesseract.image_to_string(image, lang=tesseract_lang)
        
        # 清理临时文件
        os.remove(temp_path)
        
        return {"text": text}
        
    except Exception as e:
        # 清理临时文件
        if os.path.exists(temp_path):
            os.remove(temp_path)
        raise HTTPException(status_code=500, detail=str(e))

@app.on_event("shutdown")
async def cleanup():
    """清理临时文件目录"""
    if os.path.exists(TEMP_DIR):
        shutil.rmtree(TEMP_DIR)

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8000) 